import arrow
from flask import g
from app.db.models.tasks import Task


def get_ava_results(
    task_type, limit=10, uid=None, hours=-12, ex_query=None, abnormal=False
):
    end = arrow.now()
    begin = end.shift(hours=hours)
    if not uid:
        uid = g.user.id
    if not ex_query:
        ex_query = {}
    pipeline = [
        {
            "$match": {
                "status": "completed",
                "taskType": f"{task_type}",
                "uid": uid,
                "endTime": {"$gt": begin.datetime, "$lt": end.datetime},
                **ex_query,
            }
        },
        {"$sort": {"endTime": -1}},
        {
            "$group": {
                "_id": "$taskId",
                "total": {"$sum": "$result.addition.nodeTotalCount"},
                "abnormal": {"$sum": "$result.addition.nodeAbnormalCount"},
                "jobId": {"$first": "$jobId"},
                "name": {"$first": "$name"},
                "target": {"$first": "$target"},
            }
        },
        {
            "$group": {
                "_id": "$jobId",
                "total": {"$sum": "$total"},
                "abnormal": {"$sum": "$abnormal"},
                "jobId": {"$first": "$jobId"},
                "name": {"$first": "$name"},
                "target": {"$first": "$target"},
                "taskId": {"$first": "_id"},
            }
        },
        {
            "$lookup": {
                "from": "jobs",
                "localField": "jobId",
                "foreignField": "_id",
                "as": "job",
            }
        },
    ]
    data = Task.objects.aggregate(pipeline)

    new_data = []
    for _data in data:
        if _job := _data.get("job"):
            _job = _job[0]
            _temp_dict = {
                "job_id": str(_job["_id"]),
                "name": "unknown",
                "target": "unknown",
                "availability": 0.0,
                "total": _data.get("total", 0),
                "abnormal": _data.get("abnormal", 0),
            }
            if (_total := _data.get("total")) and _total != 0:
                _temp_dict["availability"] = round(
                    (1 - (_data.get("abnormal", 0) / _total)) * 100, 2
                )
                _temp_dict["name"] = _data.get("name")
                _temp_dict["target"] = _data.get("target")
            if abnormal and _temp_dict["availability"] < 100:
                new_data.append(_temp_dict)
            elif not abnormal:
                new_data.append(_temp_dict)
    new_data.sort(key=lambda x: x["availability"])
    return new_data[:limit]


def get_ipv6_ava_results(uid, job_id):
    data_list = get_ava_results(
        "http",
        limit=1,
        uid=uid,
        hours=-24,
        ex_query={"jobId": job_id, "sourceIp": None},
    )
    if data_list:
        return data_list[0]
    return {}
